package cn.itcast.dmp.untils

import com.typesafe.config.{Config, ConfigFactory}
import org.apache.kudu.client.CreateTableOptions
import org.apache.kudu.spark.kudu.KuduContext
import org.apache.spark.internal.{Logging, config}
import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
import org.apache.spark.sql.types.StructType

class KuduUtils extends Logging {
  // 加载Kudu配置文件
  val config: Config = ConfigFactory.load("kudu.conf")

  // 由于创建Kudu表使用KuduContext对象
  var spark: SparkSession = _
  var kuduContext: KuduContext = _

  // DataFrame变量
  var dataframe: DataFrame = _

  // 定义附属构造方法，初始化SparkSession和KuduContext实例对象
  def this(sparkSession: SparkSession){
    this() // 调用主构造方法
    this.spark = sparkSession
    this.kuduContext = new KuduContext(
      config.getString("kudu.master"), // 设置Kudu Master地址
      sparkSession.sparkContext
    )
  }

  // 定义附属构造方法，传递DataFrame实例对象，初始化
  def this(df: DataFrame){
    this(df.sparkSession)  // 调用其他附属构造方法
    this.dataframe = df
  }

  /**
    * 创建Kudu中的表，当表存在时且允许删除时，先删除再创建
    *
    * @param tableName 表的名称
    * @param schema 表的Schema信息
    * @param keys 表的主键
    * @param isDelete 表若存在是否允许删除
    */
  def createKuduTable(tableName: String, schema: StructType,
                      keys: Seq[String], isDelete: Boolean = true): Unit ={

    // 判断要创建的表是否存在，如果存在，并且允许删除，则先删除再创建；否则不创建
    if(kuduContext.tableExists(tableName)){
      if(isDelete){
        kuduContext.deleteTable(tableName)
        logInfo(s"要创建的表：$tableName 已存在Kudu中，此时已经删除.....")
      }else{
        logInfo("要创建的表：$tableName 已存在Kudu中，不允许删除.....")
        return
      }
    }
    // 设置Kudu表的分区策略和副本数
    val options: CreateTableOptions = new CreateTableOptions()

    // TODO：设置主键字段作为哈希分区字段
    import scala.collection.JavaConverters._
    options.addHashPartitions(keys.asJava, 3)

    // TODO: 设置副本数
    options.setNumReplicas(config.getInt("kudu.table.factor"))

    val kuduTable = kuduContext.createTable(tableName, schema, keys, options)
    logInfo(s"创建Kudu表：$tableName(${kuduTable.getTableId}) 成功。。。。。。。。。。。。")
  }


  /**
    * 删除Kudu中存在的表
    *
    * @param tableName 表的名称
    */
  def deleteKuduTable(tableName: String): Unit ={
    // 判断表是否存在，存在就删除
    if(kuduContext.tableExists(tableName)){
      val response = kuduContext.deleteTable(tableName)
      logInfo(s"Kudu中表: $tableName 已被删除，耗时 = ${response.getElapsedMillis}......")
    }
  }

  /**
    * 将DataFrame数据保存到Kudu表中
    *
    * @param tableName 表的名称
    */
  def saveAsKuduTable(tableName: String): Unit ={
    // 保证程序健壮性，判断DataFrame是否不为null
    if(null == dataframe){
      logInfo("保存DataFrame数据为Null，重新构建新的DataFrame进行保存数据")
    }else{
      // 导入隐式转换函数
      import org.apache.kudu.spark.kudu.KuduDataFrameWriter
      dataframe.write
        .mode(SaveMode.Append)
        // 设置Kudu Master地址
        .option("kudu.master", config.getString("kudu.master"))
        // 设置读取Kudu表的名称
        .option("kudu.table", tableName)
        .kudu
    }
  }

  /**
    * SparkSession从Kudu表中读取数据
    *
    * @param tableName 表的名称
    * @return
    */
  def readKuduTable(tableName: String): Option[DataFrame] = {
    // 当Kudu表不存在时，返回None
    if(kuduContext.tableExists(tableName)){
      import org.apache.kudu.spark.kudu.KuduDataFrameReader
      val kuduDF = spark.read
        // 设置Kudu Master地址
        .option("kudu.master", config.getString("kudu.master"))
        // 设置读取Kudu表的名称
        .option("kudu.table", tableName)
        .kudu
      Some(kuduDF)
    }else{
      None
    }
  }
}


object KuduUtils extends Logging{

  // TODO: 隐式转换 -> 将SparkSession转换为KuduUtils实例对象
  implicit def sparkSessionToKuduUtils(spark: SparkSession): KuduUtils = {
    new KuduUtils(spark)
  }

  // TODO: 隐式转换 -> 将DataFrame转换为KuduUtils实例对象
  implicit def dataframeToKuduUtils(dataframe: DataFrame): KuduUtils = {
    new KuduUtils(dataframe)
  }

}

